import os,cv2,h5py,glob,copy
import numpy as np

def compare(x,y):
     stat_x = os.path.basename(x)
     stat_y = os.path.basename(y)
     if stat_x < stat_y:
          return -1
     elif stat_x > stat_y:
          return 1
     else:
          return 0

def get_data(dataset):
    images_dir = os.path.join(os.getcwd(), dataset, 'images')
    labels_dir = os.path.join(os.getcwd(), dataset, 'labels')
    images = glob.glob(os.path.join(images_dir, '*.png'))
    labels = glob.glob(os.path.join(labels_dir, '*.png'))
    images.sort()
    labels.sort()
    return images, labels

def prepare_data(dataset_path,img_w=512,img_h=512):
    images = []
    labels = []
    trimaps = []
    images_dir,labels_dir = get_data(dataset_path)
    num_img = len(images_dir)
    num_label = len(labels_dir)
    if num_img != num_label:
       raise Exception('The amounts of Images and Labels are inequal')
    else:
       num = num_label
    for i in range (num):
       image_dir = images_dir[i]
       label_dir = labels_dir[i]
       image = cv2.imread(image_dir)
       label = cv2.imread(label_dir)
       label = cv2.resize(label,(img_w,img_h))
       if len(label.shape)>=3:
           label=label[:,:,0]
       lh,lw = label.shape
       label=np.reshape(label,(lh,lw,1))
       ih,iw,ic = image.shape
       if (lw,lh) !=  (iw,ih):
           image =cv2.resize(image,(lw,lh))
    
       images.append(image)
       labels.append(label)
       trimap = label.copy()
       for i in range (lw):
           for j  in range (lh):
               if label[i,j,0]!=0 and label[i,j,0]!=255:
                  trimap[i,j,0]=128 
       trimaps.append(trimap)
    return np.array(images), np.array(labels), np.array(trimaps)

def process_data(images_dir,labels_dir,img_w=512,img_h=512):
    images = []
    labels = []
    trimaps = []
    num_img = len(images_dir)
    num_label = len(labels_dir)
    if num_img != num_label:
       raise Exception('The amounts of Images and Labels are inequal')
    else:
       num = num_label
    for i in range (num):
       image_dir = images_dir[i]
       label_dir = labels_dir[i]
       image = cv2.imread(image_dir)
       label = cv2.imread(label_dir)
       label = cv2.resize(label,(img_w,img_h))
       if len(label.shape)>=3:
           label=label[:,:,0]
       lh,lw = label.shape
       label=np.reshape(label,(lh,lw,1))
       ih,iw,ic = image.shape
       if (lw,lh) !=  (iw,ih):
           image =cv2.resize(image,(lw,lh))
    
       images.append(image)
       labels.append(label)
       trimap = label.copy()
       for i in range (lw):
           for j  in range (lh):
               if label[i,j,0]!=0 and label[i,j,0]!=255:
                  trimap[i,j,0]=128 
       trimaps.append(trimap)
    return np.array(images), np.array(labels), np.array(trimaps)


   
       #cv2.imshow('Image',image)
       #print('image shape :',image.shape)
       #cv2.waitKey(0)
       #cv2.imshow('Trimap',trimap)
       #print('Trimap shape :',trimap.shape)
       #cv2.waitKey(0)
       #cv2.imshow('Label',label)
       #print('Label shape :',label.shape)
       #cv2.waitKey(0)
       #cv2.destroyAllWindows()
       
        


if __name__=='__main__':
   image = cv2.imread('./data/images/GT01.png')
   #image=cv2.resize(image,(512,1024))
   print(image.shape)
   #prepare_data('data')
  
        
        
